#!/usr/bin/env python
# encoding: utf-8
"""
irk.py

Implicit Runge-Kutta methods in ODEPY.

ODEPY is a Python module which
consists of a library of time stepping integrators,
as well as utility functions. It is meant to facilitate the
construction and analysis of numerical integrators for 
initial value problems.

The module requires a full installation of Numpy/Scipy.

The code is available under GNU GPL v3 licence at
http://code.google.com/p/odepy

Created by Klas Modin on 2011-04-15.
"""

from numpy import array, dot, ones, zeros, zeros_like, hstack, kron, sqrt
from numpy import diag, identity, isscalar, linspace, shape, arange, finfo
from scipy.linalg import lu_solve, lu_factor, pinv
from numpy.linalg import inv, norm
from odepy.core import OneStepMethod

class ImplicitRungeKutta(OneStepMethod):
	"""
	Class for implicit Runge-Kutta methods.
	
	Attributes
	----------
	A : ndarray, shape (s,s)
		The A matrix for the method, where s is the number of stages.
	
	b : ndarray, shape (s,)
		The b vector for the method.
	
	f : function
		The RHS of the equation. Should have the form
			"def f(y)"
		and should return an ndarray of same shape as y.
		The input y should not be overwritten by f.
	
	jac : function
		The Jacobian corresponding to the RHS. Should
		have the form
			"def jac(y)"
		and should return an ndarray one order higher
		than that of y. The input y should not be
		overwritten by jac.
	
	params : dict
		Parameters for tolerances etc.
	
	stats : dict
		Statistics about total number of steps, Jacobians,
		rhs evaluations etc.
	
	Methods
	-------
	integrate
		Carry out the integration process.
	"""
	
	name = "IRK"
	
	def __init__(self, f, A, b, jac=None, params=dict()):
		"""
		Constructor.
		
		Parameters
		----------
		f : function
			The RHS of the equation. Should have the form
				"def f(y)"
			and should return an ndarray of same shape as y.
		
		A : ndarray, shape (s,s)
			The A matrix in the Butcher tableau for the method, 
			where s is the number of stages.
			Matrix should be strictly lower triangular 
			(so that the method is explicit).
		
		b : ndarray, shape (s,)
			The b vector in the Butcher tableau for the method.
		
		jac : function (optional)
			The Jacobian corresponding to the RHS. Should
			have the form
				"def jac(y)"
			and should return an ndarray one order higher
			than that of y. The input y should not be
			overwritten by jac. If not jac is given, a
			finite difference approximation will be used.
		
		params : dict (optional)
			Parameters for the root solving:
			- newtol (real), Newton tolerance
			- kmax (int), maximal number of iterations
			- kappa (real), safety factor
		"""
		super(ImplicitRungeKutta, self).__init__()
		dims = A.shape
		if (dims[0] != 1 and dims[0] != dims[1]):
			raise AttributeError("A must be a square matrix.")
		s = dims[0]
		self.A = array(A)
		if (len(b) != s):
			raise AttributeError("Size of matrix A and vector b does not match.")
		self.b = array(b)
		self.f = f
		if jac is None:
			# Notice: this Jacobian only works in the case when y is
			# a one-dimensional array.
			def jac(y):
				f0 = f(y)
				# ynew = y.copy()
				delta = 1e-10
				# for i,v in ndenumerate(ynew):
				# 	ynew(i) += delta
				J = identity(len(y))
				for Jk in J:
					Jk[:] = (f(y+delta*Jk)-f0)/delta
				return J.T
			
		self.jac = jac
		
		# Setup parameters
		self.params = dict()
		try:
			self.params['newtol'] = params['newtol']
		except KeyError:
			self.params['newtol'] = sqrt(finfo(1.0).eps) # Sqrt of machine epsilon (for default float type).
		try:
			self.params['kmax'] = params['kmax']
		except KeyError:
			self.params['kmax'] = 7
		try:
			self.params['kappa'] = params['kappa']
		except KeyError:
			self.params['kappa'] = 1e-1
		
		# Setup stats
		self.stats = dict()
		self.stats['njac'] = 0
		self.stats['nlu'] = 0
		self.stats['nrhs'] = 0
		self.stats['nsteps'] = 0
		self.stats['niterations'] = 0
		
		
	
	def setup(self,y):
		# Dimension variables
		s = len(self.b)
		n = len(y)
		
		self.hlast = 0
		self.updateJacobian = True
		self.indices = arange(s*n).reshape((s,n))
		self.Z = zeros(s*n)
		self.dZ = zeros_like(self.Z)
		self.eta = 0
		self.Uround = finfo(0.0).eps
		if (det(self.A) != 0.0):
			self.d = dot(self.b,inv(self.A))
	
	def __call__(self,h,y):
		
		if (hasattr(self,'hlast') == False):
			self.setup(y)
		
		# Dimension variables
		s = len(self.b)
		n = len(y)
		
		# Check what needs to be updated:
		if (self.hlast != h):
			self.updateLU = True
			self.hlast = h
		
		if self.updateJacobian:
			self.J = self.jac(y)
			self.updateLU = True
			self.stats['njac'] += 1
		
		if self.updateLU:
			(self.lu,self.piv) = lu_factor(identity(s*n)-h*kron(self.A,self.J))
			self.stats['nlu'] += 1
			self.updateLU = False
		
		# Set initial Z
		self.Z[:] = 0.0
		
		# First evaluation of F at Z
		self.F = kron(ones(s),self.f(y))
		self.stats['nrhs'] += 1
		
		# Iterate
		for k in xrange(self.params['kmax']):
			
			# Update statistics
			self.stats['niterations'] += 1
			
			# One Newton step
			self.dZ = lu_solve((self.lu,self.piv),h*dot(kron(self.A,identity(n)),self.F)-self.Z)
			
			# Update Z
			self.Z += self.dZ
			
			# Check stopping criterion
			if (k==0):
				self.eta = max(self.eta,self.Uround)**0.8
				self.dZnorm = norm(self.dZ)
			else:
				dZnormlast = self.dZnorm
				self.dZnorm = norm(self.dZ)
				Theta = self.dZnorm/dZnormlast
				if Theta > 1:
					raise AssertionError("The iteration is diverging with Theta=%f"%Theta)
				self.eta = Theta/(1-Theta)
			if (self.eta*self.dZnorm <= self.params['kappa']*self.params['newtol']):
				break
			
			# Evaluate F
			for i in self.indices:
				self.F[i] = self.f(y+self.Z[i])
				self.stats['nrhs'] += 1
			
		else:
			raise AssertionError("Iteration did not converge in %i iterations."%(k+1))
		
		# Update phase space variable
		if (hasattr(self,'d')):
			y += dot(self.d,self.Z[self.indices])
		else:
			# (I-h*AoJ)(Z1-Z0) = -Z0 + h AoI F(Z0)
			# Z1 = h*AoI F(Z1)
			# y1 = y0 + h*dot(b,F(Z1))
			
			# Evaluate F
			for i in self.indices:
				self.F[i] = self.f(y+self.Z[i])
				self.stats['nrhs'] += 1
			y += h*dot(self.b,self.F[self.indices])
			
		
		# Set what to do next step
		if (k==0 or Theta < 1e-3):
			self.updateJacobian = False
		else:
			self.updateJacobian = True
		
		# Update statistics
		self.stats['nsteps'] += 1
		self.stats['iterations/step'] = float(self.stats['niterations'])/self.stats['nsteps']
		self.stats['jac/step'] = float(self.stats['njac'])/self.stats['nsteps']
		self.stats['rhs/step'] = float(self.stats['nrhs'])/self.stats['nsteps']
		
		# Return output
		return y
	

class BackwardEuler(ImplicitRungeKutta):
	"""
	The backward (implicit) Euler method. 
	"""
	name = "Backward Euler"
	def __init__(self,f,*args,**kwargs):
		print args
		print kwargs
		A = array([[1.0]])
		b = array([1.0])
		super(BackwardEuler, self).__init__(f,A,b,*args,**kwargs)
	

class ImplicitMidpoint(ImplicitRungeKutta):
	"""
	The implicit midpoint method (2nd order Gauss). 
	This is a symplectic Runge-Kutta method.
	"""
	name = "Implicit midpoint"
	def __init__(self,f,*args,**kwargs):
		A = array([[0.5]])
		b = array([1.0])
		super(ImplicitMidpoint, self).__init__(f,A,b,*args,**kwargs)
	

class Gauss4(ImplicitRungeKutta):
	"""
	The 4th order Gauss method. 
	This is a symplectic Runge-Kutta method.
	"""
	name = "Gauss-4"
	def __init__(self,f,*args,**kwargs):
		A = array([[0.25,0.25-sqrt(3.0)/6],[0.25+sqrt(3.0)/6,0.25]])
		b = array([0.5,0.5])
		super(Gauss4, self).__init__(f,A,b,*args,**kwargs)
	

class Gauss6(ImplicitRungeKutta):
	"""
	The 6th order Gauss method. 
	This is a symplectic Runge-Kutta method.
	"""
	name = "Gauss-6"
	def __init__(self,f,*args,**kwargs):
		A = array([
			[5.0/36 , 2.0/9-sqrt(15.0)/15 , 5.0/36-sqrt(15.0)/30],
			[5.0/36+sqrt(15.0)/24 , 2.0/9 , 5.0/36-sqrt(15.0)/24],
			[5.0/36+sqrt(15.0)/30 , 2.0/9+sqrt(15.0)/15 , 5.0/36]
		])
		b = array([5.0/18,4.0/9,5.0/18])
		super(Gauss6, self).__init__(f,A,b,*args,**kwargs)
	

class RadauIA3(ImplicitRungeKutta):
	"""
	The 3rd order Radau IA method. 
	"""
	name = "Radau IA-3"
	def __init__(self,f,*args,**kwargs):
		A = array([
			[0.25 , -0.25],
			[0.25 , 5.0/12]
		])
		b = array([0.25,0.75])
		super(RadauIA3, self).__init__(f,A,b,*args,**kwargs)
	

class RadauIA5(ImplicitRungeKutta):
	"""
	The 5th order Radau IA method. 
	"""
	name = "Radau IA-5"
	def __init__(self,f,*args,**kwargs):
		A = array([
			[1.0/9 , (-1.0-sqrt(6.0))/18 , (-1.0+sqrt(6.0))/18],
			[1.0/9 , (88.0+7*sqrt(6.0))/360 , (88.0-43*sqrt(6.0))/360],
			[1.0/9 , (88.0+43*sqrt(6.0))/360 , (88.0-7*sqrt(6.0))/360]
		])
		b = array([1.0/9,(16.0+sqrt(6.0))/36,(16.0-sqrt(6.0))/36])
		super(RadauIA5, self).__init__(f,A,b,*args,**kwargs)
	

class RadauIIA3(ImplicitRungeKutta):
	"""
	The 3rd order Radau IIA method. 
	"""
	name = "Radau IIA-3"
	def __init__(self,f,*args,**kwargs):
		A = array([
			[5.0/12 , -1.0/12],
			[0.75 , 0.25]
		])
		b = array([0.75,0.25])
		super(RadauIIA3, self).__init__(f,A,b,*args,**kwargs)
	

class RadauIIA5(ImplicitRungeKutta):
	"""
	The 5th order Radau IIA method. 
	"""
	name = "Radau IIA-5"
	def __init__(self,f,*args,**kwargs):
		A = array([
			[(88.0-7*sqrt(6.0))/360 , (296.0-169*sqrt(6.0))/1800 , (-2.0+3*sqrt(6.0))/225],
			[(296.0+169*sqrt(6.0))/1800 , (88.0+7*sqrt(6.0))/360 , (-2.0-3*sqrt(6.0))/225],
			[(16.0-sqrt(6.0))/36 , (16.0+sqrt(6.0))/36, 1.0/9]
		])
		b = array([(16.0-sqrt(6.0))/36 , (16.0+sqrt(6.0))/36, 1.0/9])
		super(RadauIIA5, self).__init__(f,A,b,*args,**kwargs)
	


if __name__ == '__main__' and __file__.split('/')[-1] == 'irk.py':
	from numpy import hstack, sin, cos, pi
	# from pylab import plot,semilogy,clf,legend
	execfile('odepy/irk_list.py')
	
	# Setup non-linear pendulum problem
	y0 = hstack((pi*0.99,0))
	def f(y):
		dim = len(y)/2
		q = y[:dim]
		p = y[dim:]
		return hstack((p,-sin(q)))
		# return hstack((p,-q))
		# return hstack((p,-q))
	
	def jac(y):
		dim = len(y)/2
		q = y[:dim]
		p = y[dim:]
		return array([[0,1.0],[-cos(q),0]])
		# return array([[0,1.0],[-1.0,0]])
		# return array([[0,1.0],[-1.0,0]])
	
	def H(y):
		q = y[:,0]
		p = y[:,1]
		return p*p/2-cos(q)
	
	H0 = H(array([y0]))
	
	solverlist = []
	params = dict()
	params['newtol'] = 1e-10
	# solverlist.append(BackwardEuler(f,jac))
	# solverlist.append(ImplicitMidpoint(f,jac,params))
	# solverlist.append(RadauIA3(f,jac))
	# solverlist.append(RadauIIA3(f,jac,params))
	# solverlist.append(Gauss_4(f,jac,params))
	# solverlist.append(RadauIA5(f,jac,params))
	# solverlist.append(RadauIIA5(f,jac))
	# solverlist.append(Gauss_6(f,jac))
	solverlist.append(LobattoIIIA_2(f,jac,params))
	solverlist.append(LobattoIIIB_2(f,jac,params))
	# solverlist.append(LobattoIIIC_2(f,jac,params))
	# solverlist.append(LobattoIIICstar_2(f,jac,params))
	# solverlist.append(LobattoIIID_2(f,jac,params))
	# solverlist.append(LobattoIIIA_2(f,jac,params)**0.5*LobattoIIIB_2(f,jac,params)**0.5)
	# solverlist.append(LobattoIIIB_4(f,jac,params))
	# solverlist.append(LobattoIIIB_6(f,jac,params))
	# solverlist.append(LobattoIIIB_8(f,jac,params))
	
	
	tend = 100.0
	stepsize = 0.3
	Houtlist = []
	for solver in solverlist:
		(tout,yout) = solver.integrate(tend,y0,stepsize)
		Houtlist.append(abs(H(yout)-H0))
	Hout = array(Houtlist).T
	
	# clf();semilogy(tout,Hout)
	# legend([str(s) for s in solverlist],loc='best')


